import networkx as nx
import pulp
import itertools
from itertools import product

# Symbolic expression classes for order of magnitude estimates (where we only care about magnitudes up to a constant factor)


# Statements are formal assertions involving expressions X, Y, such as X \lesssim Y, X \sim Y, or X \gtrsim Y.
# 
class Statement:
    def __init__(self, left, op, right):
        self.left = left
        self.op = op
        self.right = right
    def __str__(self):
        return f"{self.left} {self.op} {self.right}"
    def __repr__(self):
        return str(self)

# An expression is a formal combination of variables, where the variables represent orders of magnitude, combined using sums, products, quotients, powers, min, and max.

class Expression:
    def __add__(self, other):
        return Add(self, ensure_expr(other))
    def __radd__(self, other):
        return Add(ensure_expr(other), self)
    def __mul__(self, other):
        return Mul(self, ensure_expr(other))
    def __rmul__(self, other):
        return Mul(ensure_expr(other), self)
    def __truediv__(self, other):
        return Div(self, ensure_expr(other))
    def __rtruediv__(self, other):
        return Div(ensure_expr(other), self)
    def __pow__(self, other):
        return Power(self, ensure_expr(other))
    def __rpow__(self, other):
        return Power(ensure_expr(other), self)
    def __hash__(self):
        # Use the object’s identity so that two different objects
        # are always distinct in hash tables.
        return id(self)
    def __le__(self, other):
        return Statement(self, '≲', ensure_expr(other))
    def __ge__(self, other):
        return Statement(self, '≳', ensure_expr(other))
    def __eq__(self, other):  # need equality for hash comparison; use comparable() instead for X ~ Y
        return self is other  
    def __str__(self):
        raise NotImplementedError("Must implement __str__ in subclasses")
    def __repr__(self):
        return str(self)
    def clean(self): # treats all constants as a single object `one` (in particular, all constants become equal) 
        if isinstance(self, Constant):
            return one
        else:
            return self

# Variable is our class for non-negative quantities, defined up to an order of magnitude
class 变量(Expression):
    def __init__(self, name):
        self.name = name
    def __str__(self):
        return self.name

class Add(Expression):
    def __init__(self, left, right):
        self.left = left
        self.right = right
    def __str__(self):
        # Parenthesize to preserve clarity in nested expressions
        return f"({self.left} + {self.right})"

class Mul(Expression):
    def __init__(self, left, right):
        self.left = left
        self.right = right
    def __str__(self):
        # Parenthesize to preserve clarity in nested expressions
        return f"({self.left} * {self.right})"

class Div(Expression):
    def __init__(self, left, right):
        self.left = left
        self.right = right
    def __str__(self):
        # Parenthesize to preserve clarity in nested expressions
        return f"({self.left} / {self.right})"

class Power(Expression):
# TODO: typecheck that exponent is a constant
    def __init__(self, base, exponent):
        self.base = base
        self.exponent = exponent
    def __str__(self):
        # Parenthesize to preserve clarity in nested expressions
        return f"({self.base} ** {self.exponent})"

class Max(Expression):
    def __init__(self, *operands):
        self.operands = operands
    def __str__(self):
        inner = ", ".join(str(op) for op in self.operands)
        return f"max({inner})"

class Constant(Expression):
    def __init__(self, value):
        if not isinstance(value, (int, float)):
            raise TypeError("Constant only accepts int/float")
        self.value = value

    def __str__(self):
        # drop the .0 for ints
        if isinstance(self.value, int) or self.value.is_integer():
            return str(int(self.value))
        return str(self.value)

one = Constant(1)  # need a global constant here for equality checks (other instances of Constant(1) might not be equal to one.  Perhaps there is a more pythonic way to do this?)

class Min(Expression):
    def __init__(self, *operands):
        self.operands = operands
    def __str__(self):
        inner = ", ".join(str(op) for op in self.operands)
        return f"min({inner})"

def ensure_expr(obj):
    """If obj is already an Expression, return it;
       otherwise wrap it in a Constant."""
    return obj if isinstance(obj, Expression) else Constant(obj)

# Currently, assumptions are in two types: `assumptions`, which are `statements`, and `splits`, which are disjunctions of conjunctions of statements.  One could conceivably add more types of assumptions.
class 假设:
    def __init__(self):
        self.assumptions = []
        self.splits = []
    
    def add(self, statement):
        print(f"Adding assumption: {statement}")
        self.assumptions.append(statement)

    def add_lp(self, *lp_statement):
        print(f"Adding Littlewood-Paley assumption that expressions of magnitudes {lp_statement} can sum to zero.")
        self.splits.append(cases_lp(lp_statement))

    def split_at(self, expr1, expr2):
        print(f"Force a case split at {expr1} == {expr2}.")
        self.splits.append([[expr1 <= expr2], [expr2 <= expr1]])

    def 可证(self, expr1, expr2):
        """Checks if a statement can be proven from the assumptions."""
        print(f"基于给定公理，确定 {expr2} 是否 {expr1} 的上界。")

        # start collecting all the case splittings we will use.  First we collect the expressions that would require splittings appearing in the bound to be proven, as well as the ambient assumptions.
        base_splits = splittings(expr1 <= expr2)
        for assumption in self.assumptions:
            base_splits.update(splittings(assumption))
        
        # For each such expression (which may be a max or min), we collect the cases that would be required to prove the bound.
        cases = []
        for split in base_splits:
            if isinstance(split, Max):
                cases.append(cases_max(*split.operands))
            elif isinstance(split, Min):
                cases.append(cases_min(*split.operands))

        # we also allow for the assumptions to contain some forced splittings.
        for split in self.splits:
            cases.append(split)

        if len(cases) > 0:
            print("分为以下情况：")
            for case in cases:
                print(case)

        for axioms in product(*cases):
            if len(cases) > 0:
                print(f"\n检查情况: {axioms}")
            ordering = 排序有向图()
            for ax in axioms:
                for statement in ax:
                    ordering.add(statement)
# TODO: simplify the hypotheses before adding them to the ordering
            for assumption in self.assumptions:
                ordering.add(assumption)
            if not ordering.可证(expr1, expr2):
                return False

        print("所有情况皆得证!")
        return True

# An Ordering directed graph, where an edge from A to B means A is bounded by a constant times B
class 排序有向图:
    def __init__(self):
        self.graph = nx.DiGraph()
        self.statements = set()

    def __str__(self):
        return(str(self.statements))
         
    """Assumes a statement, adding it to the directed graph of of hypotheses."""
    def add(self, statement):
        self.statements.add(statement)
        match statement.op:
            case '≲':
                self.graph.add_edge(statement.left.clean(), statement.right.clean())
            case '≳':
                self.graph.add_edge(statement.right.clean(), statement.left.clean())
            case '~':
                self.graph.add_edge(statement.left.clean(), statement.right.clean())
                self.graph.add_edge(statement.right.clean(), statement.left.clean())
            case _:
                raise ValueError(f"Unknown operator: {statement.op}")

# sees if a statement can be directly proven from the hypotheses and transitivity
    def 有通路(self, statement):
        match statement.op:
            case '≲':
                return nx.has_path(self.graph, statement.left.clean(), statement.right.clean())
            case '≳':
                return nx.has_path(self.graph, statement.right.clean(), statement.left.clean())
            case '~':
                return (nx.has_path(self.graph, statement.left.clean(), statement.right.clean()) and
                        nx.has_path(self.graph, statement.right.clean(), statement.left.clean())) 

    def maximal_elements(self, expressions):
        """Returns the potentially maximal elements in a set of expressions."""
        new_expressions = expressions.copy()
        for expr1, expr2 in itertools.combinations(expressions, 2):
            if expr1 in new_expressions and expr2 in new_expressions:
                if self.有通路(expr1 <= expr2):
                    new_expressions.discard(expr1)
                elif self.有通路(expr2 <= expr1):
                    new_expressions.discard(expr2)
        return new_expressions
    
    def minimal_elements(self, expressions):
        """Returns the potentially minimal elements in a set of expressions."""
        new_expressions = expressions.copy()
        for expr1, expr2 in itertools.combinations(expressions, 2):
            if expr1 in new_expressions and expr2 in new_expressions:
                if self.有通路(expr1 <= expr2):
                    new_expressions.discard(expr2)
                elif self.有通路(expr2 <= expr1):
                    new_expressions.discard(expr1)
        return new_expressions

    def 排序简化(self, expr):
        """Simplifies an expression using the ordering hypotheses, e.g., simplifying max(X,Y) to Y if we know X ≲ Y."""
#        print(f"Attempting to simplify {expr} using the ordering hypotheses {self}.")
        if isinstance(expr, 变量):
            return expr
        elif isinstance(expr, Constant):
            return one
        elif isinstance(expr, Add):
            return self.排序简化(max(expr.left, expr.right))  # outsource processing of Add to max  
        elif isinstance(expr, Mul):
# TODO: collect any Mul or Prod instances in the operands, and simplify them to a single Prod instance
            left = self.排序简化(expr.left)
            right = self.排序简化(expr.right)
            if isinstance(left, Constant):
                return right
            if isinstance(right, Constant):
                return left
            return Mul(left, right)
        elif isinstance(expr, Div):
            left = self.排序简化(expr.left)
            right = self.排序简化(expr.right)
            if isinstance(right, Constant):
                return left
            return Div(left, right)
        elif isinstance(expr, Power):
            base = self.排序简化(expr.base)
            if isinstance(base, Constant):
                return one
            return Power(base, expr.exponent)
        elif isinstance(expr, Max):
# TODO: concatenate any maxima in the operands to create a single maximum
            new_operands = self.maximal_elements( {self.排序简化(op) for op in expr.operands} )
            if len(new_operands) == 0:
                return one
            if len(new_operands) == 1:
                return new_operands.pop()
            return Max(*new_operands)
        elif isinstance(expr, Min):
# TODO: concatenate any minima in the operands to create a single minimum
            new_operands = self.minimal_elements( {self.排序简化(op) for op in expr.operands} )
            if len(new_operands) == 0:
                return one
            if len(new_operands) == 1:
                return new_operands.pop()
            return Min(*new_operands)
        elif isinstance(expr, Statement):
            left = self.排序简化(expr.left)
            right = self.排序简化(expr.right)
            return Statement(left, expr.op, right)
        elif isinstance(expr, (int, float)):
            return one
        else:
            raise TypeError(f"Unsupported expression type: {type(expr)}")

# tests if one can prove expr1 ≲ expr2 using the hypotheses
    def 可证(self, expr1, expr2):
#        print(f"Attempting to bound {expr1} by {expr2} using {self}.")
        expr = self.排序简化(expr2/expr1)
        print(f"化为证明 {monomial_simplify(expr)} >= 1.")
        terms = monomials(expr)
        for key, _ in terms.items():
            self.graph.add_node(key)  # Ensure all variables are nodes in the graph

        prob = pulp.LpProblem("BoundCheck", pulp.LpMinimize)
                
        # Create variables
        w = { (u,v): pulp.LpVariable(f"w_{u}_{v}", lowBound=0) for u, v in self.graph.edges() }

        # Dummy objective: minimize total weight
        prob += pulp.lpSum(w)
        
        # Constraints for each term in the expression, except one
        for v in self.graph.nodes():
            if not v == one:
                inflow = pulp.lpSum(w[u, v] for u, v in self.graph.in_edges(v))
                outflow = pulp.lpSum(w[v, u] for v, u in self.graph.out_edges(v))
                prob += (inflow - outflow == terms.get(v, 0), f"FlowBalance_{v}")
        
        silent_solver = pulp.PULP_CBC_CMD(msg=False)
        prob.solve(silent_solver)
        if pulp.LpStatus[prob.status] == 'Optimal':
            if pulp.value(prob.objective) > 0:
                print(f"渐进界由以下假设相乘得证：")
                for (u,v), var in w.items():
                    if var.value() > 0:
                        print(f"{u} ≲ {v} 的 {var.value()} 次幂")
            else:
                print("显然。")
            return True
        else:  
            print("无法证明为界。")
            return False

# Collect all the terms in a given expression that require splitting into cases
def splittings(expr):
    if isinstance(expr, (变量, Constant)):
        return set()
    elif isinstance(expr, Add):
        splits = splittings(expr.left).union(splittings(expr.right))
        splits.add(Max(expr.left, expr.right))
        return splits
    elif isinstance(expr, (Mul, Div)):
        return splittings(expr.left).union(splittings(expr.right))
    elif isinstance(expr, Power):
        return splittings(expr.base).union(splittings(expr.exponent))
    elif isinstance(expr, (Max,Min)):
        splits = set()
        for operand in expr.operands:
            splits.update(splittings(operand))
        splits.add(expr)
        return splits
    elif isinstance(expr, Statement):
        left_splits = splittings(expr.left)
        right_splits = splittings(expr.right)
        return left_splits.union(right_splits)
    elif isinstance(expr, (int, float)):
        return set()
    else:
        raise TypeError(f"Unsupported expression type: {type(expr)}")

# define all the splittings associated to max(args)
def cases_max(*args):
    cases = []
    if len(args) < 2:
        return cases
    for expr in args:
        ordering = []
        for other_expr in args:
            if expr != other_expr:
                ordering.append(other_expr <= expr)
        cases.append(ordering)
    return cases

# define all splittings associated to min(args)
def cases_min(*args):
    cases = []
    if len(args) < 2:
        return cases
    for expr in args:
        ordering = []
        for other_expr in args:
            if expr != other_expr:
                ordering.append(expr <= other_expr)
        cases.append(ordering)
    return cases

# define all splittings associated to a Littlewood-Paley constraint (sum to zero).  In such a context, two of the variables are comparable, and the others are bounded by the first two.
def cases_lp(args):
    cases = []
    assert len(args) == 1, "Error: Littlewood-Paley constraints cannot involve just one variable."
    for expr1, expr2 in itertools.combinations(args, 2):
        ordering = [comparable(expr1,expr2)]
        for other_expr in args:
            if other_expr != expr1 and other_expr != expr2:
                ordering.append(other_expr <= expr1)
        cases.append(ordering)
    return cases

# return the simplified expression formed by gathering terms
def monomial_simplify(expr):
    if isinstance(expr, Statement):
        # Simplify the left and right expressions of the statement
        left = monomial_simplify(expr.left)
        right = monomial_simplify(expr.right)
        return Statement(left, expr.op, right)
    else:
        first = True
        result = one
        for key, value in monomials(expr).items():
            if value == 0:
                factor = one
            elif value == 1:
                factor = key
            else:
                factor = key ** value
            if first:
                result = factor
                first = False
            else:
                result = result * factor
        return result  

# return a dictionary where the keys are the terms in the expression and the values are the exponents in which they appear
def monomials(expr):
    """Returns a dictionary of monomials in the expression with their exponents."""
    if isinstance(expr, (变量, Add, Max, Min)):
        return {expr: 1}
    elif isinstance(expr, Constant):
        return {}  # ignore all constants
    elif isinstance(expr, Mul):
        left_monomials = monomials(expr.left)
        right_monomials = monomials(expr.right)
        # Multiply the dictionaries (add exponents for common terms)
        for key, value in right_monomials.items():
            left_monomials[key] = left_monomials.get(key, 0) + value
        return left_monomials
    elif isinstance(expr, Div):
        numerator = monomials(expr.left)
        denominator = monomials(expr.right)
        # Subtract exponents for division
        for key, value in denominator.items():
            numerator[key] = numerator.get(key, 0) - value
        return numerator
    elif isinstance(expr, Power):
        base_monomials = monomials(expr.base)
        exponent = expr.exponent.value if isinstance(expr.exponent, Constant) else 1  # Assume exponent is a constant for now
        # Raise each term to the power of the exponent
        for key in base_monomials.keys():
            base_monomials[key] *= exponent
        return base_monomials
    else:
        raise TypeError(f"Unsupported expression type: {type(expr)}")

# we cannot use `expr1 == expr2` to denote comparability, because we need to reserve the equality symbol for several other purposes.
def comparable(expr1, expr2):
    """Returns a statement that compares two expressions."""
    if isinstance(expr1, Expression) and isinstance(expr2, Expression):
        return Statement(expr1, '~', expr2)
    else:
        raise TypeError("Both arguments must be Expression instances")
